package org.loong.crypto.extension.cms.jce;

import java.security.PrivateKey;
import java.security.interfaces.ECPrivateKey;
import java.security.interfaces.RSAPrivateKey;
import java.util.Objects;

import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x9.X9ObjectIdentifiers;
import org.bouncycastle.operator.AsymmetricKeyUnwrapper;
import org.bouncycastle.operator.GenericKey;
import org.bouncycastle.operator.OperatorException;
import org.loong.crypto.common.exception.CryptoException;
import org.loong.crypto.core.algorithm.CipherAlgorithm;
import org.loong.crypto.core.algorithm.EncryptionMethod;
import org.loong.crypto.core.params.CipherParameters;
import org.loong.crypto.extension.operator.SM2Utils;
import org.loong.crypto.model.asn1.sm2.SM2Cipher;
import org.loong.crypto.service.CryptoService;

public class JceAsymmetricKeyUnwrapper extends AsymmetricKeyUnwrapper {

    private CryptoService cryptoService;

    private PrivateKey privateKey;

    public JceAsymmetricKeyUnwrapper(AlgorithmIdentifier encAlgId, CryptoService cryptoService, PrivateKey privateKey) {
        super(encAlgId);

        this.cryptoService = cryptoService;
        this.privateKey = privateKey;
    }

    /**
     * Generate unwrapped key.
     * 
     * @param encryptedKeyAlgorithm the encryptedKeyAlgorithm
     * @param encryptedKey the encryptedKey
     * @return the {@link GenericKey}
     * @throws OperatorException {@link OperatorException}
     */
    public GenericKey generateUnwrappedKey(AlgorithmIdentifier encryptedKeyAlgorithm, byte[] encryptedKey) throws OperatorException {
        if (Objects.isNull(privateKey)) {
            throw new OperatorException("PrivateKey cannot be null.");
        }

        try {
            CipherAlgorithm algorithm = null;
            CipherParameters.Builder paramsBuilder = CipherParameters.builder().key(privateKey);
            ASN1ObjectIdentifier algorithmId = getAlgorithmIdentifier().getAlgorithm();
            if (algorithmId.equals(PKCSObjectIdentifiers.rsaEncryption)) {
                algorithm = CipherAlgorithm.RSA;
                paramsBuilder.encryptionMethod(EncryptionMethod.RSA_ECB_PKCS1);
            } else if (algorithmId.equals(X9ObjectIdentifiers.id_ecPublicKey)) {
                algorithm = CipherAlgorithm.SM2;
            } else {
                throw new OperatorException("Unsupported algorithm: " + algorithmId);
            }

            byte[] key = null;
            if (privateKey instanceof RSAPrivateKey) {
                key = cryptoService.decrypt(algorithm, paramsBuilder.build(), encryptedKey);
            } else if (privateKey instanceof ECPrivateKey) {
                SM2Cipher sm2Cipher = SM2Cipher.getInstance(encryptedKey);
                if (sm2Cipher == null) {
                    throw new OperatorException("Unable to recover sm2Cipher.");
                }

                byte[] encryptedData = SM2Utils.reduceCipherText(sm2Cipher);
                key = cryptoService.decrypt(algorithm, paramsBuilder.build(), encryptedData);
            } else {
                throw new OperatorException("Unknown private key type.");
            }

            if (Objects.isNull(key)) {
                throw new OperatorException("Unable to recover secret key.");
            }

            return new GenericKey(encryptedKeyAlgorithm, key);
        } catch (CryptoException e) {
            throw new OperatorException(e.getMessage(), e);
        }
    }
}
